use super::*;

use crate::bin_tree::BinMultTree;
use crate::bin_var_node::{CtrlMsg, CtrlMsgA};
use crate::check_msg::CheckMsg;

use belief_propagation::{BPError, BPResult, NodeFunction, NodeIndex};
use rustfft::{num_complex::Complex, Fft};
use std::sync::Arc;
use std::convert::TryInto;

pub struct CheckNode<const K: usize, const ETA: usize> {
    iter_step: i16, //---迭代步数------
    eqindex: i16, //--方程编号------------
    narrowed_times: i16, //--已经收缩的次数--
    thresholds: [f64; 21], //--阈值，最多收缩10轮--
    sweep_efforts: [i16; 21], //--清扫力度----
    n: usize,
    coeffs: [i16; K],
    value: [i16; 3], //--------现在value上放了3个值---------------------
    op: CmpOperator,
    fft: Arc<dyn Fft<f64>>,
    ifft: Arc<dyn Fft<f64>>,
    connections: Vec<usize>,
}

impl<const K: usize, const ETA: usize> CheckNode<K, ETA> {
    pub fn new(iter_step: i16, eqindex: i16, narrowed_times: i16, thresholds: [f64; 21], sweep_efforts: [i16; 21], coeffs: [i16; K], value: [i16; 3], op: CmpOperator, n: usize) -> Self { //------现在value上放了3个值---添加了方程编好eqIndex--添加了迭代步数iter_step--添加了收缩次数narrowed_times---------
        let mut planner = FftPlanner::new();
        let fft = planner.plan_fft_forward(n);
        let ifft = planner.plan_fft_inverse(n);
        Self {
            iter_step: iter_step,
            eqindex: eqindex,
            narrowed_times: narrowed_times,
            thresholds: thresholds,
            sweep_efforts: sweep_efforts,
            n: n,
            coeffs: coeffs,
            value: value,
            op: op,
            connections: Vec::with_capacity(K),
            fft,
            ifft,
        }
    }
    fn entropy(&self, data: &[f64; ETA]) -> f64 {
        -data
            .iter()
            .map(|p| {
                let r = p * p.log2();
                if r.is_nan() {
                    0.0
                } else {
                    r
                }
            })
            .sum::<f64>()
    }
    //--------计算是否和 均匀分布 比较像的函数---------------------------------------------
    fn js_divergence(&self,
        data: &[f64; ETA]  
    ) -> f64 {
        //因为data中的分布所做的正规化是除以最大值的正规化，本身构不成一个分布，所以需要转化成概率分布再比较
        //println!("input data {:?}", data.to_vec());

        let summ = {
            data.iter()
                .sum::<f64>()
        };
        if summ == 0.0 {
            return 0.0
        } 
        let norm_data : [f64; ETA] = {
            data.iter()
                .map(|d| {
                    d/summ
                })
                .collect::<Vec<f64>>()
                .try_into()
                .unwrap()
        };
        //println!("after normalize {:?}", norm_data.to_vec());
        let m_dist : [f64; ETA] = {
            norm_data.iter()
                .map(|d| {
                   0.5 * (d + 1.0/(ETA as f64))
                })
                .collect::<Vec<f64>>()
                .try_into()
                .unwrap()
        };
        //println!("mix dist {:?}", m_dist.to_vec());
        let js : f64 = self.entropy( &m_dist) - 0.5 * (self.entropy(&norm_data) + (ETA as f64).log2());
        if js < 0.0 {
            println!("Negetive JS Divergence: {:?}", data.to_vec());
        }
        //println!("js = {}", js);
        js.sqrt()  // 平方根，是js距，而非js散度
    }
    //-----计算当前方程在各系数下js均值指标，即直接用于比较的指标----------------------
    fn extract_indicator(&self, js_div: &Vec<(NodeIndex, f64)>) -> f64 {
        //每个系数类别下先求均值
        let mut sum_coeff: [f64; 9] = [0.0; 9]; 
        let mut count_coeff: [i32; 9] = [0; 9];
        
        for (nodeindex, js_value) in js_div {
            sum_coeff[(self.coeffs[*nodeindex]+ 4) as usize ] += js_value;
            count_coeff[(self.coeffs[*nodeindex]+ 4) as usize ] += 1;
        }
        for x in 0..count_coeff.len() {
            if count_coeff[x]==0 {
                count_coeff[x] = 1;
            }
        }
        let mean_coeff : [f64; 9] = {sum_coeff
            .iter()
            .zip(count_coeff.iter())
            .map(|(s,c)|  *s / *c as f64)
            .collect::<Vec<f64>>()
            .try_into()
            .unwrap()
        };
        //非零的系数类别下的值再求均值。改为-3，-2，-1,1,2,3六个系数类别下求均值
        let average = (mean_coeff[1] +mean_coeff[2] + mean_coeff[3] + mean_coeff[5] + mean_coeff[6] + mean_coeff[7] ) / 6 as f64;
        average
    }
    //----------------------------------------------------------------
    fn node_function_normal(
        & mut self, //使得不等式节点CheckNode的边界可改变，加上了mut
        inbox: Vec<(NodeIndex, CheckMsg<ETA>)>,
    ) -> BPResult<Vec<(NodeIndex, CheckMsg<ETA>)>> {
        //println!("{:?}", inbox[0].1.data.to_vec());
        let leafs: Vec<Vec<Complex<f64>>> = inbox
            .iter()
            .map(|(node_index, msg)| {
                msg.mult_and_transform(self.coeffs[*node_index], self.n, &self.fft)
                    .iter()
                    .map(|p| p / (self.n as f64).sqrt())
                    .collect()
            })
            .collect();

        let products: Vec<Vec<Complex<f64>>> =
            BinMultTree::new(leafs, multiply_pointwise).calculate();

        let partials: Vec<Vec<f64>> = products 
            .into_iter()
            .map(|prd| ifft(prd, &self.ifft))
            .collect();
        
        //if self.eqindex == 0 {
         //   println!("step = {}", self.iter_step);
         //   println!("Sum Dist: {:?}", partials[0].clone().into_iter());
        //}
        
        
        //-----------------计算边缘分布--------------------------------------
           
        if self.iter_step <=6 {   
        //if self.iter_step <= -1 {    //原版，不收缩
            if  self.narrowed_times == 0  { //如果是未收缩过的方程
            
                let margin_res: Vec<(NodeIndex, CheckMsg<ETA>)> = match self.op{
                    //------For double bounds--------------------------
                    CmpOperator::DoubleEq => partials.clone()
                        .into_iter()
                        .zip(inbox.clone().into_iter())
                        .map(|(dist_sum, ib)|{
                            (
                                ib.0,
                                derive_margin_from_inequality_double_eq(dist_sum, self.value[0], self.value[1], self.coeffs[ib.0]),
                            )
                        })
                        .collect(),
                    //--注意单边的还没添加margin分布的函数，这里只是保持代码结构完整----------
                    CmpOperator::GreaterEq => partials.clone()
                        .into_iter()
                        .zip(inbox.clone().into_iter())
                        .map(|(dist_sum, ib)|{
                            (
                                ib.0,
                                //todo
                                derive_margin_from_inequality_double_eq(dist_sum, self.value[0], self.value[1], self.coeffs[ib.0]),
                            )
                        })
                        .collect(),
                    //--注意单边的还没添加margin分布的函数，这里只是保持代码结构完整----------
                    CmpOperator::SmallerEq => partials.clone()
                        .into_iter()
                        .zip(inbox.clone().into_iter())
                        .map(|(dist_sum, ib)|{
                            (
                                ib.0,
                                //todo
                                derive_margin_from_inequality_double_eq(dist_sum, self.value[0], self.value[1], self.coeffs[ib.0]),
                            )
                        })
                        .collect(),
                };
                //-------------------------------------------------------------------------------------------
                let margin_res_1 = margin_res.to_vec().clone();
                //判断下边缘分布中的概率是否有负数，如果有，说明线性组合的真实值的概率还没有凸显出来，还没有到可以收缩的时候
                let mut mature_eq: i16 = 1; //是否是已经成熟的方程
                for (_, margin_dis) in margin_res_1.clone().into_iter() {
                    for pr in margin_dis.data {
                        if pr < 0.0 {
                            mature_eq = 0; //但凡出现负的概率表明还不够成熟，不能收缩
                            break; 
                        }
                    }
                    if mature_eq == 0 {
                        break;
                    }
                }
                //---判断margin_res和均匀分布之间的距离是否足够大-------
                
                //边缘分布中可能有负数（因为涉及的线性组合的分布中已经有负数了），认为这种情况下真实值距离边缘很远，直接需要收缩，
                //----partials中存放的是当前方程中除去一个节点后剩余的节点的线性组合的分布，怀疑其中有负数---------
                /* 
                for comb_dist in partials.clone() {
                    
                    for pr in &comb_dist {
                        if *pr < 0.0 {
                            println!("{:.4?}, ", comb_dist);//看是不是负数都分布在两侧很远的地方，还是都是负数
                        
                            
                            break;
                        }
                    
                    }
                    
                }//--------------------
                */

                if mature_eq == 1 { // 只有成熟的方程，有资格收缩的方程才输出indicator值参与阈值的决策
                    let js_div: Vec<(NodeIndex, f64)> = {
                        margin_res_1.into_iter()
                        .map(|(nodeindex, res_dist)| {
                            (
                                nodeindex,
                                self.js_divergence(&res_dist.data),
                            )        
                        })
                        .collect()
                    };
                    //这里是输出（每个方程、每个未知数、对应的JSD）三元组，方便进行初步分析，从图上判断从哪轮开始收缩更有效，额外预分析用什么指标更有效
                    for (_,js) in js_div.iter().enumerate()
                    {   
                        println!("{}, {}, {}", self.eqindex, js.0, js.1);            
                    };


                    /* 
                    let indicator: f64 = self.extract_indicator(&js_div);
                    //println!("{}, {}", self.eqindex, indicator); 
                    
                    if indicator < self.thresholds[self.iter_step as usize] { //如果小于阈值且是成熟的方程，考虑收缩
                        
                        //if self.value[0] + self.sweep_efforts[self.iter_step as usize] < self.value[1] - self.sweep_efforts[self.iter_step as usize] {
                            //println!("narrowing, {}, {}", self.eqindex,  self.sweep_efforts[self.iter_step as usize]);
                            self.value[0] += self.sweep_efforts[self.iter_step as usize];
                            self.value[1] -= self.sweep_efforts[self.iter_step as usize];
                            self.narrowed_times += 1;
                        //}
                        
                    }
                    */
                }      
            } 
            
            /* 
            if self.iter_step == 4 { //收缩完了全部初始化为最初的分布，self.iter_step加1，重新迭代
                
                let mut prior_msg = CheckMsg::new();
                
                
                for v in -(ETA as i16) / 2..=(ETA as i16) / 2 {
                    /* 
                    if v >= 0 {
                        prior_msg[v] = (0 .. ((ETA as u64) / 2 + 1)).map(|x| binomial.pmf(x) * binomial.pmf(x + v as u64)).sum();
                        
                    } else {
                        prior_msg[v] = (0 .. ((ETA as u64) / 2 + 1)).map(|x| binomial.pmf(x) * binomial.pmf(x - v as u64)).sum();
                    }  
                    */
                    prior_msg[v] = - 1.0;  //传回奇异的负数，这时变量节点那边就更改为初始值        
                }
                //println!("Dist: {:?}", prior_msg.data.to_vec());
                //prior_msg.normalize().expect("Failed to normalize re-initialed dist");
                
                //println!("Dist: {:?}", prior_msg.data.to_vec());
                
                let res4: Vec<(NodeIndex, CheckMsg<ETA>)> = inbox.into_iter()
                                                            .map(|ib|{
                                                                (
                                                                    ib.0,
                                                                    prior_msg,
                                                                )
                                                            })
                                                            .collect();

                
                //let res4: Vec<(NodeIndex, CheckMsg<ETA>)> = Vec::with_capacity(0);
                self.iter_step +=1;
                return Ok(res4);

            }
            */
        }
        //----------------------------------------------------------------------------
        
        let res: Vec<(NodeIndex, CheckMsg<ETA>)> = match self.op {
            //--------For double bounds------------------------------------------------------------------------------------
            CmpOperator::DoubleEq => partials
                .into_iter()
                .zip(inbox.into_iter())
                .map(|(dist_sum, ib)| {
                    (
                        ib.0,
                        derive_from_inequality_double_eq(dist_sum, self.value[0], self.value[1], self.coeffs[ib.0]),
                    )
                })
                .collect(),
            //---------------------------------------------------------------------------------------------------------------
            CmpOperator::GreaterEq => partials
                .into_iter()
                .zip(inbox.into_iter())
                .map(|(dist_sum, ib)| {
                    (
                        ib.0,
                        derive_from_inequality_greater(dist_sum, self.value[0], self.coeffs[ib.0]), //注意单边不等式时界都放在了value[0]里
                    )
                })
                .collect(),

            CmpOperator::SmallerEq => partials
                .into_iter()
                .zip(inbox.into_iter())
                .map(|(dist_sum, ib)| {
                    (
                        ib.0,
                        derive_from_inequality_smaller(dist_sum, self.value[0], self.coeffs[ib.0]), //注意单边不等式时界都放在了value[0]里
                    )
                })
                .collect(),
        };
        
        
        //----------------------------------------------------------------------------
        self.iter_step +=1;
        //if (self.eqindex == 0) {
        //    println!("iter_step = {}", self.iter_step);
        //}
        Ok(res)
    }
}

impl<const K: usize, const ETA: usize> NodeFunction<i16, CheckMsg<ETA>, CtrlMsg, CtrlMsgA>
    for CheckNode<K, ETA>
{
    fn node_function(
        &mut self,
        inbox: Vec<(NodeIndex, CheckMsg<ETA>)>,
    ) -> BPResult<Vec<(NodeIndex, CheckMsg<ETA>)>> {
        self.node_function_normal(inbox)
    }

    fn number_inputs(&self) -> Option<usize> {
        Some(K)
    }

    fn is_factor(&self) -> bool {
        true
    }
    fn get_prior(&self) -> Option<CheckMsg<ETA>> {
        None
    }
    fn initialize(&mut self, connections: Vec<NodeIndex>) -> BPResult<()> {
        //TODO: Ensure connections are sorted
        if connections.len() != K {
            Err(BPError::new(
                "CheckNode::initialize".to_owned(),
                format!(
                    "Wrong number ({}) of connections given ({}).",
                    K,
                    connections.len()
                ),
            ))
        } else {
            self.connections = connections;
            Ok(())
        }
    }
    fn reset(&mut self) -> BPResult<()> {
        Ok(())
    }

    fn is_ready(
        &self,
        recv_from: &Vec<(NodeIndex, CheckMsg<ETA>)>,
        _current_step: usize,
    ) -> BPResult<bool> {
        Ok(recv_from.len() == self.connections.len())
    }
}
